clear
close all

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Parameters
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
mu      = 0;                    %mean of log(theta)
sigma   = 0.175;                %sd of log(theta)
xmin    = 0.01;                 %lowest value of theta grid
xmax    = 100.01;               %highest value of theta grid
N       = 100001;               %number of grid points
beta    = 66/76;                %:=1/alpha
delta   = 0.9432;               %see section 5.4 for choice

ExpTheta = exp(mu+0.5*sigma^2); %expectation of theta
T=200;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Get Cutoffs
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
[SpendRes,SpendIrr,DurRes,DurIrr,Prob,ExpSpend,~,~,~,~,~,thetae,~,~,~,thetastar,thetahat,star,starstar,nstar,nstarstar] = Sim(mu,sigma,xmin,xmax,N,beta,delta);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Simulate theta
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

rng(7);
for i=1:T
    theta(i) = lognrnd(mu,sigma);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% a) First-Best (Spending without bias)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

gfirst = theta./(theta+delta/(1-delta)*ExpTheta);

Bfirst(1) = 0.2;

for i=1:T
    Bfirst(i+1) = 1/delta*(1-gfirst(i))*Bfirst(i)+(1-1/delta*(1-gfirst(i)))/(1-delta);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% b) Markov (Continuously flexible spending)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

gmarkov = theta./(theta+beta*delta/(1-delta)*ExpTheta);

Bmarkov(1) = 0.2;

for i=1:T
    Bmarkov(i+1) = 1/delta*(1-gmarkov(i))*Bmarkov(i)+(1-1/delta*(1-gmarkov(i)))/(1-delta);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% c) PPE (Best equilibrium)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

R(1) = 1;
for i=1:T
    if R(i) == 1
        if theta(i)<star
            gppe(i) = theta(i)./(theta(i)+beta*delta/(1-delta)*ExpTheta);
            R(i+1) = 1;
        end
        if star <= theta(i) && theta(i) <= starstar
            gppe(i) = star./(star+beta*delta/(1-delta)*ExpTheta);
            R(i+1) = 1;
        end
        if theta(i)>starstar
            gppe(i) = theta(i)./(theta(i)+beta*delta/(1-delta)*ExpTheta);
            R(i+1) = 0;
        end
    end
    if R(i)==0
        if theta(i)<nstarstar
            gppe(i) = theta(i)./(theta(i)+beta*delta/(1-delta)*ExpTheta);
            R(i+1) = 0;
        end
        if nstarstar <= theta(i) && theta(i)  <= nstar
            gppe(i) = nstar./(nstar+beta*delta/(1-delta)*ExpTheta);
            R(i+1) = 1;
        end
        if theta(i)>nstar
            gppe(i) = theta(i)./(theta(i)+beta*delta/(1-delta)*ExpTheta);
            R(i+1) = 1;
        end
    end
end

Bppe(1) = 0.2;

for i=1:T
    Bppe(i+1) = 1/delta*(1-gppe(i))*Bppe(i)+(1-1/delta*(1-gppe(i)))/(1-delta);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Plot
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Black regions (Flexible rate exceeding threshold under Fiscal Resp.
aboveThreshold = (gmarkov > star./(star+beta*delta/(1-delta)*ExpTheta)).*R(2:201).*R(1:200);
startIndex = find(diff([false, aboveThreshold]) == 1);
endIndex = find(diff([aboveThreshold, false]) == -1);
startIndex = max(startIndex - 1 , 1); 
endIndex = min(endIndex , T); 

figure(1)
axes('FontSize', 24, 'NextPlot', 'add');
hold on
p1 = plot([1:T],gfirst,':','color',[0,0,0],'linewidth',1.5)
p2 = plot([1:T],gmarkov,'--','color',[0,0,0],'linewidth',1.5)
p3 = plot([1:T],gppe,'-','color',[0,0,0],'linewidth',1.5), xlabel('Time','FontSize',24), ylabel('Spending Rate','FontSize',24)
yline(star./(star+beta*delta/(1-delta)*ExpTheta),'-','color',[0.5,0.5,0.5])
yline(nstar./(nstar+beta*delta/(1-delta)*ExpTheta),'-','color',[0.5,0.5,0.5])
hold off
for i=1:T
    if R(i)==0
        patch('XData',[i-1 i-1 i i],'YData', [0 0.12 0.12 0],'EdgeColor','none','FaceColor','k','FaceAlpha',0.1)
    end
end
for i = 1:length(startIndex)
    xStart = 1+startIndex(i) - 0.2;  
    xEnd = endIndex(i) + 0.2;      
    width = xEnd - xStart;    
    rectangle('Position', [xStart, 0, width, .003], 'FaceColor', 'k', 'EdgeColor', 'none');
end
xlim([0 80])
ylim([0 .12])
legend([p1,p2,p3],{'First Best','Flexible','Best Equilibrium'},'FontSize',24)
set(gcf, 'Units', 'Normalized', 'OuterPosition', [.04, .08, .90, .80]);
saveas(gcf,'SpendingRates.png')
